#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import time
import sys
import paramiko
import re

def print_step(message):
    """
    打印带有格式的步骤信息
    """
    print(f"\n{'=' * 60}")
    print(f"  {message}")
    print(f"{'=' * 60}\n")

def execute_command(ssh, command, print_output=True):
    """
    在SSH会话中执行命令并返回输出
    """
    print(f"执行命令: {command}")
    stdin, stdout, stderr = ssh.exec_command(command)
    exit_status = stdout.channel.recv_exit_status()
    
    output = stdout.read().decode('utf-8')
    error = stderr.read().decode('utf-8')
    
    if print_output:
        if output:
            print(f"输出:\n{output}")
        if error:
            print(f"错误:\n{error}")
    
    return output, error, exit_status

def wait_for_ssh(host, port=22, timeout=300):
    """
    等待SSH服务可用
    """
    import socket
    start_time = time.time()
    while True:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(5)
            sock.connect((host, port))
            sock.close()
            return True
        except socket.error:
            if time.time() - start_time > timeout:
                return False
            time.sleep(5)

def check_if_installed(ssh, package_name):
    """
    检查软件包是否已安装
    """
    if package_name == "docker":
        output, _, _ = execute_command(ssh, "docker --version", print_output=False)
        return "Docker version" in output
    elif package_name == "wireguard":
        output, _, _ = execute_command(ssh, "docker ps -a | grep wg-easy", print_output=False)
        return "wg-easy" in output
    return False

def install_docker(ssh):
    """
    安装Docker
    """
    print_step("检查Docker是否已安装")
    if check_if_installed(ssh, "docker"):
        print("Docker已经安装，跳过安装步骤")
        return True
    
    print_step("更新软件包列表")
    execute_command(ssh, "sudo apt-get update -y")
    
    print_step("安装Docker依赖")
    execute_command(ssh, "sudo apt-get install -y apt-transport-https ca-certificates curl software-properties-common")
    
    print_step("添加Docker官方GPG密钥")
    execute_command(ssh, "curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -")
    
    print_step("添加Docker仓库")
    execute_command(ssh, 'sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"')
    
    print_step("再次更新软件包列表")
    execute_command(ssh, "sudo apt-get update -y")
    
    print_step("安装Docker CE")
    execute_command(ssh, "sudo apt-get install -y docker-ce docker-ce-cli containerd.io")
    
    print_step("启动Docker服务")
    execute_command(ssh, "sudo systemctl start docker")
    execute_command(ssh, "sudo systemctl enable docker")
    
    print_step("验证Docker安装")
    output, error, status = execute_command(ssh, "docker --version")
    if status != 0:
        print("Docker安装失败")
        return False
    
    print_step("将当前用户添加到docker组")
    execute_command(ssh, "sudo usermod -aG docker $USER")
    
    return True

def install_wireguard(ssh, server_ip):
    """
    通过Docker安装WireGuard (wg-easy)
    """
    print_step("检查WireGuard是否已安装")
    if check_if_installed(ssh, "wireguard"):
        print("WireGuard已经安装，跳过安装步骤")
        return True
    
    print_step("创建WireGuard配置目录")
    execute_command(ssh, "mkdir -p ~/.wg-easy")
    
    print_step("拉取WireGuard (wg-easy) Docker镜像")
    execute_command(ssh, "sudo docker pull ghcr.io/wg-easy/wg-easy")
    
    print_step("启动WireGuard (wg-easy) 容器")
    wireguard_command = (
        "sudo docker run -d "
        "--name=wg-easy "
        "-e LANG=en "
        f"-e WG_HOST={server_ip} "
        "-e PASSWORD_HASH='$2a$12$Wc/NFt21TXolJ1Dyw7GMW.Ln.P7Jw6F1RUD5a9AXdxqwRwpGfWFmW' "
        "-e PORT=51821 "
        "-e WG_PORT=51820 "
        "-e WG_DEFAULT_DNS=8.8.4.4,1.1.1.1 "
        "-v ~/.wg-easy:/etc/wireguard "
        "-p 51820:51820/udp "
        "-p 10086:51821/tcp "
        "--cap-add=NET_ADMIN "
        "--cap-add=SYS_MODULE "
        "--sysctl=\"net.ipv4.conf.all.src_valid_mark=1\" "
        "--sysctl=\"net.ipv4.ip_forward=1\" "
        "--restart unless-stopped "
        "ghcr.io/wg-easy/wg-easy"
    )
    execute_command(ssh, wireguard_command)
    
    print_step("等待WireGuard容器启动")
    time.sleep(10)
    
    print_step("检查WireGuard容器状态")
    output, _, _ = execute_command(ssh, "sudo docker ps | grep wg-easy")
    if "wg-easy" not in output:
        print("WireGuard容器启动失败")
        return False
    
    print_step("WireGuard Web UI信息")
    print(f"WireGuard Web UI地址: http://{server_ip}:10086")
    print("默认密码已设置，请通过Web UI管理WireGuard")
    
    return True

def main():
    parser = argparse.ArgumentParser(description='安装Docker和WireGuard到远程服务器')
    parser.add_argument('ip', help='远程服务器IP地址')
    parser.add_argument('username', help='SSH用户名')
    parser.add_argument('password', help='SSH密码')
    parser.add_argument('--port', type=int, default=22, help='SSH端口 (默认: 22)')
    
    args = parser.parse_args()
    
    print_step(f"连接到远程服务器 {args.ip}")
    
    try:
        # 创建SSH客户端连接
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(args.ip, port=args.port, username=args.username, password=args.password)
        
        print("SSH连接成功")
        
        # 安装Docker
        if not install_docker(ssh):
            print("Docker安装失败，退出程序")
            ssh.close()
            return 1
        
        # 安装WireGuard，使用IP地址作为WG_HOST参数
        if not install_wireguard(ssh, args.ip):
            print("WireGuard安装失败，退出程序")
            ssh.close()
            return 1
        
        print_step("安装完成")
        print("Docker和WireGuard (wg-easy) 已成功安装！")
        print(f"WireGuard Web UI地址: http://{args.ip}:10086")
        print("您可以通过Web UI管理WireGuard，添加和配置客户端")
        
        # 关闭SSH连接
        ssh.close()
        return 0
        
    except paramiko.AuthenticationException:
        print("认证失败，请检查用户名和密码")
        return 1
    except paramiko.SSHException as e:
        print(f"SSH连接错误: {str(e)}")
        return 1
    except Exception as e:
        print(f"发生错误: {str(e)}")
        return 1

if __name__ == "__main__":
    sys.exit(main())